package main

import (
	"bytes"
	"crypto/md5"
	"crypto/sha1"
	"errors"
	"hash"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
)

var xJar = XJar{
	md5:  []byte{#{xJar.md5}},
	sha1: []byte{#{xJar.sha1}},
}

var xKey = XKey{
	algorithm: []byte{#{xKey.algorithm}},
	keysize:   []byte{#{xKey.keysize}},
	ivsize:    []byte{#{xKey.ivsize}},
	password:  []byte{#{xKey.password}},
}

func main() {
	jar, err := JAR(os.Args)
	if err != nil {
		panic(err)
	}

	{
		java := os.Args[1]
		javaPath := ""
		if strings.HasPrefix(java, "/") {
			javaPath = java
		} else {
			s, err := existJava()
			if err != nil {
				panic(err)
			}
			javaPath = s
		}

		s1, err := SHA1(javaPath)
		if err != nil {
			panic(err)
		}

		// https://download.oracle.com/otn/java/jdk/8u201-b09/42970487e3af4f5aa5bca3f542482c60/jdk-8u201-linux-x64.rpm?AuthParam=1665560257_0099839d2edaf45b1c0418f590785cc2
		// java
		s2 := []byte{110, 82, 39, 1, 145, 176, 7, 39, 250, 229, 28, 136, 52, 201, 187, 36, 181, 130, 188, 15}
		if bytes.Compare(s1, s2) != 0 {
			panic(errors.New("invalid java with hash"))
		}

		path, err := filepath.Abs(jar)
		if err != nil {
			panic(err)
		}
		SHA1, err := SHA1(path)
		if err != nil {
			panic(err)
		}
		if bytes.Compare(SHA1, xJar.sha1) != 0 {
			panic(errors.New("invalid jar with hash"))
		}
	}

	myArgs := make([]string, 0)

	{
		attachMechanism := false

		args := os.Args
		l := len(args)
		for i := 0; i < l; i++ {
			arg := args[i]
			if strings.HasPrefix(arg, "-javaagent:") {
				panic(errors.New("agent forbidden"))
			}

			if strings.Contains(arg, "-XX:+TraceClassLoading") ||
				strings.Contains(arg, "-XX:+CompileTheWorld") ||
				strings.Contains(arg, "-XX:+CompileTheWorldPreloadClasses") {
				panic(errors.New(arg + " forbidden"))
			}
			if strings.EqualFold(arg, "-XX:+DisableAttachMechanism") {
				attachMechanism = true
			}
		}

		if !attachMechanism {
			myArgs = append(myArgs, "-XX:+DisableAttachMechanism")
		}
	}

	// start java application
	java := os.Args[1]
	args := append(myArgs, os.Args[2:]...)
	key := bytes.Join([][]byte{
		xKey.algorithm, {13, 10},
		xKey.keysize, {13, 10},
		xKey.ivsize, {13, 10},
		xKey.password, {13, 10},
	}, []byte{})
	cmd := exec.Command(java, args...)
	cmd.Stdin = bytes.NewReader(key)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	err = cmd.Run()
	if err != nil {
		panic(err)
	}
}

func MD5(path string) ([]byte, error) {
	return HASH(path, md5.New())
}

func SHA1(path string) ([]byte, error) {
	return HASH(path, sha1.New())
}

func HASH(path string, hash hash.Hash) ([]byte, error) {
	file, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	_, _err := io.Copy(hash, file)
	if _err != nil {
		return nil, _err
	}

	sum := hash.Sum(nil)
	return sum, nil
}


// find jar name from args
func JAR(args []string) (string, error) {
	var jar string

	l := len(args)
	for i := 1; i < l-1; i++ {
		arg := args[i]
		if arg == "-jar" {
			jar = args[i+1]
		}
	}

	if jar == "" {
		return "", errors.New("unspecified jar name")
	}

	return jar, nil
}

type XJar struct {
	md5  []byte
	sha1 []byte
}

type XKey struct {
	algorithm []byte
	keysize   []byte
	ivsize    []byte
	password  []byte
}

// exists isExist dir or file
func exists(path string) bool {
	_, err := os.Stat(path)
	if err != nil {
		return os.IsExist(err)
	}
	return true
}

func isDir(path string) bool {
	s, err := os.Stat(path)
	if err != nil {
		return false
	}
	return s.IsDir()
}

func existJava() (string, error) {
	sysPath := os.Getenv("PATH")
	sp := strings.Split(sysPath, ":")
	for _, s := range sp {
		if !exists(s) {
			continue
		}

		if isDir(s) {
			javaPath := filepath.Join(s, "java")
			if exists(javaPath) {
				link, err := os.Readlink(javaPath)
				if err != nil {
					return javaPath, nil
				}
				return link, nil
			}
		} else {
			if strings.HasSuffix(s, "java") {
				link, err := os.Readlink(s)
				if err != nil {
					return filepath.Join(s, "java"), nil
				}
				return link, nil
			}
		}
	}

	return "", errors.New("java is not exist in PATH")
}
